package com.ruoyi.video.server.utils;

import com.ruoyi.common.config.RuoYiConfig;
import com.ruoyi.video.domain.VideoCamera;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.io.IOUtils;
import org.bytedeco.javacv.FFmpegFrameGrabber;
import org.bytedeco.javacv.FFmpegFrameRecorder;
import org.bytedeco.javacv.Frame;
import org.bytedeco.javacv.Java2DFrameConverter;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.ByteArrayOutputStream;
import java.io.File;

/**
 * RTSP流工具
 */
@Slf4j
public class StreamUtils {

    /**
     * 截图
     * @param rtsp 流地址
     * @return
     */
    public static String snap(String rtsp){
        String imgPath = RuoYiConfig.getDownloadPath();
        File imgFolder = new File(imgPath);
        imgFolder.mkdirs();

        FFmpegFrameGrabber grabber = new FFmpegFrameGrabber(rtsp);
        try {

            grabber.start(); // 初始化抓取器

            // 抓取第5秒时的一帧
            grabber.setTimestamp(3L);
            Frame frame = grabber.grabFrame();

            if (frame != null) {
                // 将帧转换为BufferedImage
                Java2DFrameConverter converter = new Java2DFrameConverter();
                BufferedImage bufferedImage = converter.convert(frame);

                // 保存为图片
                File outputFile = new File(imgPath+System.currentTimeMillis()+".jpg");
                ImageIO.write(bufferedImage, "JPG", outputFile);
                return outputFile.getName();
            }
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            if (grabber != null) {
                try {
                    grabber.stop(); // 停止抓取器
                }catch (Exception e){
                    e.printStackTrace();
                }
            }
        }
        return null;
    }

    /**
     * 测试单个RTSP流地址
     * @param rtsp
     * @return
     */
    public static boolean test(String rtsp){
        log.info("开始测试RTSP流： {}", rtsp);
        FFmpegFrameGrabber grabber = null;
        ByteArrayOutputStream stream = null;
        FFmpegFrameRecorder recorder = null;
        try {
            grabber = new FFmpegFrameGrabber(rtsp);
            if ("rtsp".equals(rtsp.substring(0, 4))) {
                grabber.setOption("rtsp_transport", "tcp");
                grabber.setOption("stimeout", "5000000");
            }
            grabber.start();
            stream = new ByteArrayOutputStream();
            recorder = new FFmpegFrameRecorder(stream, grabber.getImageWidth(), grabber.getImageHeight(),
                    grabber.getAudioChannels());
            recorder.setInterleaved(true);
            recorder.setVideoOption("preset", "ultrafast");
            recorder.setVideoOption("tune", "zerolatency");
            recorder.setVideoOption("crf", "25");
            recorder.setFrameRate(grabber.getFrameRate());
            recorder.setSampleRate(grabber.getSampleRate());
            if (grabber.getAudioChannels() > 0) {
                recorder.setAudioChannels(grabber.getAudioChannels());
                recorder.setAudioBitrate(grabber.getAudioBitrate());
                recorder.setAudioCodec(grabber.getAudioCodec());
            }
            recorder.setFormat("flv");
            recorder.setVideoBitrate(grabber.getVideoBitrate());
            recorder.setVideoCodec(grabber.getVideoCodec());
            recorder.start(grabber.getFormatContext());
            byte[] headers = stream.toByteArray();
            stream.reset();
            log.info("RTSP拉流成功： {}", rtsp);
            return true;
        } catch (Exception e) {
            e.printStackTrace();
            log.info("RTSP拉流失败： {}\r\n{}", rtsp, e);
            return false;
        } finally {
            try {
                if(grabber != null) {
                    IOUtils.close(grabber);
                }
                if(recorder != null) {
                    IOUtils.close(recorder);
                }
                if(stream != null) {
                    IOUtils.close(stream);
                }
            }catch (Exception e){
                e.printStackTrace();
            }
        }
    }
}
